package com.jackie.trpc.filter;

import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.atomic.AtomicLong;
import javax.annotation.PostConstruct;
import com.jackie.trpc.config.OpenTelemetryConfig;
import com.tencent.trpc.core.filter.spi.Filter;
import com.tencent.trpc.core.rpc.Invoker;
import com.tencent.trpc.core.rpc.Request;
import com.tencent.trpc.core.rpc.Response;
import io.opentelemetry.api.OpenTelemetry;
import io.opentelemetry.api.common.AttributeKey;
import io.opentelemetry.api.common.Attributes;
import io.opentelemetry.api.metrics.DoubleHistogram;
import io.opentelemetry.api.metrics.LongCounter;
import io.opentelemetry.api.metrics.Meter;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.SpanKind;
import io.opentelemetry.api.trace.StatusCode;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import io.opentelemetry.context.propagation.TextMapGetter;
import io.opentelemetry.context.propagation.TextMapSetter;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;


@Component
public class OpenTelemetryTracingFilter implements Filter {
    // 用于提取上下文（Server端）
    private static final TextMapGetter<Map<String, Object>> TRPC_GETTER = 
        new TextMapGetter<Map<String, Object>>() {
            @Override
            public Iterable<String> keys(Map<String, Object> carrier) {
                return carrier.keySet();
            }

            @Override
            public String get(Map<String, Object> carrier, String key) {
                Object value = carrier.get(key);
                if (value == null) return null;
                if (value instanceof byte[]) {
                    return new String((byte[]) value, StandardCharsets.UTF_8);
                }
                return value.toString();
            }
        };

    //用于注入上下文（Client端）
    private static final TextMapSetter<Map<String, Object>> TRPC_SETTER = 
        new TextMapSetter<Map<String, Object>>() {
            @Override
            public void set(Map<String, Object> carrier, String key, String value) {
                if (carrier != null && key != null && value != null) {
                    carrier.put(key, value.getBytes(StandardCharsets.UTF_8));
                }
            }
        };

    @Autowired
    private Tracer tracer;
    
    @Autowired
    private OpenTelemetryConfig config;
    
    @Autowired
    private OpenTelemetry openTelemetry;

     // 定义Metrics相关对象
    private LongCounter requestCounter;
    private LongCounter errorCounter;
    private LongCounter statusCodeCounter;
    private LongCounter timeoutCounter;
    private LongCounter slowRequestCounter;
    private LongCounter dependencyErrorCounter;
    private LongCounter retryCounter;
    private DoubleHistogram durationHistogram;
    private DoubleHistogram requestSizeHistogram;
    private DoubleHistogram responseSizeHistogram;
    private AtomicLong concurrentRequests;
    private LongCounter rateLimitedCounter;
    private LongCounter circuitBreakerOpenCounter;

    @PostConstruct
    public void initMetrics() {
        Meter meter = openTelemetry.getMeter("trpc.filter");
        // 1. 流量相关指标
        requestCounter = meter.counterBuilder("trpc.requests.total")
            .setDescription("Total number of TRPC requests")
            .build();
            
        errorCounter = meter.counterBuilder("trpc.errors.total")
            .setDescription("Total number of failed TRPC requests")
            .build();
            
        statusCodeCounter = meter.counterBuilder("trpc.requests.by_status")
            .setDescription("Requests count by status code")
            .build();
            
        concurrentRequests = new AtomicLong();
        meter.gaugeBuilder("trpc.concurrent.requests")
            .setDescription("Current concurrent requests")
            .ofLongs()
            .buildWithCallback(measurement -> 
                measurement.record(concurrentRequests.get())
            );
        // 2. 性能指标
        durationHistogram = meter.histogramBuilder("trpc.request.duration")
            .setDescription("TRPC request duration in milliseconds")
            .setUnit("ms")
            .build();
            
        timeoutCounter = meter.counterBuilder("trpc.timeouts.total")
            .setDescription("Total timed-out requests")
            .build();
            
        slowRequestCounter = meter.counterBuilder("trpc.slow_requests.total")
            .setDescription("Requests slower than threshold")
            .build();

        // 3. 资源用量指标
        requestSizeHistogram = meter.histogramBuilder("trpc.request.size.bytes")
            .setUnit("bytes")
            .setDescription("Request size distribution")
            .build();
            
        responseSizeHistogram = meter.histogramBuilder("trpc.response.size.bytes")
            .setUnit("bytes")
            .setDescription("Response size distribution")
            .build();

        // 4. 下游依赖指标
        dependencyErrorCounter = meter.counterBuilder("trpc.dependency.errors")
            .setDescription("Downstream service failures")
            .build();
            
        retryCounter = meter.counterBuilder("trpc.retries.total")
            .setDescription("Total retry attempts")
            .build();

        // 5. 系统健康度指标
        rateLimitedCounter = meter.counterBuilder("trpc.rate_limited.total")
            .setDescription("Total rate-limited requests")
            .build();
            
        circuitBreakerOpenCounter = meter.counterBuilder("trpc.circuit_breaker_open.total")
            .setDescription("Total requests rejected by circuit breaker")
            .build();
    }

     @Override
    public CompletionStage<Response> filter(Invoker<?> invoker, Request request) {
        long startTime = System.currentTimeMillis();
        concurrentRequests.incrementAndGet();
        
        // 记录请求大小
        if (request.getAttachments() != null) {
            long requestSize = request.getAttachments().values().stream()
                .mapToLong(v -> v instanceof byte[] ? ((byte[])v).length : v.toString().length())
                .sum();
            requestSizeHistogram.record(requestSize,
                Attributes.of(
                    AttributeKey.stringKey("service"), request.getInvocation().getRpcServiceName(),
                    AttributeKey.stringKey("method"), request.getInvocation().getFunc()
                ));
        }
        
        if (isOutgoingRequest(request)) {
            return handleOutgoingRequest(invoker, request, startTime);
        } else {
            return handleIncomingRequest(invoker, request, startTime);
        }
    }

    private boolean isOutgoingRequest(Request request) {
        return request.getAttachments().isEmpty();
    }

    private CompletionStage<Response> handleIncomingRequest(Invoker<?> invoker, Request request, long startTime) {
        // 基础标签
        Attributes baseAttributes = Attributes.of(
            AttributeKey.stringKey("service"), request.getInvocation().getRpcServiceName(),
            AttributeKey.stringKey("method"), request.getInvocation().getFunc()
        );
        
        // 增加请求计数
        requestCounter.add(1, baseAttributes);
            
        Map<String, Object> attachments = request.getAttachments();
        Context parentContext = openTelemetry.getPropagators().getTextMapPropagator()
            .extract(Context.current(), attachments, TRPC_GETTER);
            
        String spanName = String.format("trpc/%s/%s", 
            request.getInvocation().getRpcServiceName(),
            request.getInvocation().getFunc());
            
        Span span = tracer.spanBuilder(spanName)
            .setParent(parentContext)
            .setSpanKind(SpanKind.SERVER)
            .startSpan();
            
        try (Scope scope = span.makeCurrent()) {
            return invoker.invoke(request)
                .whenComplete((response, throwable) -> {
                    long duration = System.currentTimeMillis() - startTime;
                    concurrentRequests.decrementAndGet();
                    
                    // 记录耗时
                    Attributes durationAttributes = Attributes.builder()
                        .putAll(baseAttributes)
                        .put("status", throwable != null ? "error" : "success")
                        .build();
                    durationHistogram.record(duration, durationAttributes);
                    
                    // 记录慢请求
                    if (duration > 1000) { // 1秒阈值
                        slowRequestCounter.add(1, baseAttributes);
                    }
                    
                        // 记录响应大小
                    if (response != null && response.getValue() != null) {
                        int size = response.getValue().toString().length();
                        responseSizeHistogram.record(size, baseAttributes);
                        span.setAttribute(config.TRPC_RESPONSE_SIZE, size);
                    }
                    statusCodeCounter.add(1, 
                        Attributes.builder()
                            .putAll(baseAttributes)
                            .put("status_code", "200")
                            .build());
                    span.setStatus(StatusCode.OK);
                    span.end();
                });
        } catch (Exception e) {
            concurrentRequests.decrementAndGet();
            long duration = System.currentTimeMillis() - startTime;
            // 错误处理
            Attributes errorAttributes = Attributes.builder()
                .putAll(baseAttributes)
                .put("error.type", e.getClass().getSimpleName())
                .build();
                
            errorCounter.add(1, errorAttributes);
            statusCodeCounter.add(1, 
                Attributes.builder()
                    .putAll(baseAttributes)
                    .put("status_code", "500")
                    .build());
            
            durationHistogram.record(duration,
                Attributes.builder()
                    .putAll(baseAttributes)
                    .put("status", "error")
                    .build());
                
            span.recordException(e)
                .setStatus(StatusCode.ERROR, e.getMessage());
            span.end();
            throw e;
        }
    }


    private CompletionStage<Response> handleOutgoingRequest(Invoker<?> invoker, Request request, long startTime) {
        Attributes baseAttributes = Attributes.of(
            AttributeKey.stringKey("service"), request.getInvocation().getRpcServiceName(),
            AttributeKey.stringKey("method"), request.getInvocation().getFunc(),
            AttributeKey.stringKey("dependency"), request.getInvocation().getRpcServiceName()
        );
        
        requestCounter.add(1, baseAttributes);
            
        String spanName = String.format("trpc/%s/%s", 
            request.getInvocation().getRpcServiceName(),
            request.getInvocation().getFunc());
            
        Span span = tracer.spanBuilder(spanName)
            .setSpanKind(SpanKind.CLIENT)
            .startSpan();
            
        try (Scope scope = span.makeCurrent()) {
            openTelemetry.getPropagators().getTextMapPropagator()
                .inject(Context.current(), request.getAttachments(), TRPC_SETTER);
                
            return invoker.invoke(request)
                .whenComplete((response, throwable) -> {
                    long duration = System.currentTimeMillis() - startTime;
                    concurrentRequests.decrementAndGet();
                    
                    durationHistogram.record(duration,
                        Attributes.builder()
                            .putAll(baseAttributes)
                            .put("status", throwable != null ? "error" : "success")
                            .build());
                        
                   
                    statusCodeCounter.add(1, 
                        Attributes.builder()
                            .putAll(baseAttributes)
                            .put("status_code", "200")
                            .build());
                        span.setStatus(StatusCode.OK);
                    span.end();
                });
        } catch (Exception e) {
            concurrentRequests.decrementAndGet();
            long duration = System.currentTimeMillis() - startTime;
            
            Attributes errorAttributes = Attributes.builder()
                .putAll(baseAttributes)
                .put("error.type", e.getClass().getSimpleName())
                .build();
                
            errorCounter.add(1, errorAttributes);
            dependencyErrorCounter.add(1, baseAttributes);
            statusCodeCounter.add(1, 
                Attributes.builder()
                    .putAll(baseAttributes)
                    .put("status_code", "500")
                    .build());
            
            durationHistogram.record(duration,
                Attributes.builder()
                    .putAll(baseAttributes)
                    .put("status", "error")
                    .build());
                
            span.recordException(e)
                .setStatus(StatusCode.ERROR, e.getMessage());
            span.end();
            throw e;
        }
    }
}